import pandas as pd
import numpy as np

def cumprod_test():
    dates = pd.date_range('2016-01-01', periods=6, freq='D')
    print(dates)
    datas = np.arange(start=1.1, stop=1.7, step=0.1, dtype='float')
    print(datas)
    price = pd.DataFrame(data=datas, index=dates, columns=['p'])
    print(price)
    price['pct_change'] = price['p'].pct_change()
    print(price)

    ret_index = (1 + price['pct_change']).cumprod()
    ret_index[0] = 1
    print(ret_index)


def decile_test():
    dates = pd.date_range('2016-01-01', periods=10, freq='D')
    datas = np.arange(start=1.0, stop=2.0, step=0.1, dtype='float')
    price = pd.DataFrame(data=datas, index=dates, columns=['p'])
    print(price)
    returns = price.pct_change()
    print(returns)

    def to_index(rets):
        index = (1 + rets).cumprod()

        first_loc = max(index.notnull().argmax() - 1, 0)
        index.values[first_loc] = 1
        return index

    def trend_signal(rets, lookback, lag):
        # signal = pd.rolling_sum(rets, lookback, min_periods=lookback - 5)
        signal = rets.rolling(window=lookback, min_periods=3).sum()
        print(signal)
        return signal.shift(lag)

    signal = trend_signal(returns, 4, 3)
    print(signal)
    trade_friday = signal.resample('W-FRI').resample('B').ffill()
    print(trade_friday)
    trade_rets = trade_friday.shift(1) * returns
    print(trade_rets)

if __name__ == '__main__':
    # cumprod_test()
    decile_test()